#include "search.h"
#include "string.h"
#include <stdlib.h>
#include <fstream>
#include <chrono>
#include <thread>
#include <atomic>
#include <mutex>
#include <vector>
#include <iostream>
#include <sstream>
#include <bitset>
#include <algorithm>
#include <iomanip>
#include <cstdint>
#include "uint128_t.h"


geo_ip *ipv4InstanceLocationFile(const char *fileName)
{
	geo_ip *ipv4InstanceLocation = (geo_ip *)malloc(sizeof(geo_ip));
	if (!ipv4InstanceLocation)
	{
		std::cerr << "memory allocation failed" << std::endl;
		return NULL;
	}
	ipv4InstanceLocation->readMode = "file";
	ipv4InstanceLocation->version = 4;
	ipv4InstanceLocation->fileName = fileName;
	if (ip_loadDat_v4(ipv4InstanceLocation, fileName) >= 0)
	{
		return ipv4InstanceLocation;
	}
	free(ipv4InstanceLocation);
	ipv4InstanceLocation = NULL;
	return NULL;
}

geo_ip *ipv4InstanceLocationMemory(const char *fileName)
{
	static geo_ip *ipv4InstanceLocation = NULL;
	if (ipv4InstanceLocation == NULL)
	{
		ipv4InstanceLocation = (geo_ip *)malloc(sizeof(geo_ip));
		if (!ipv4InstanceLocation)
		{
			std::cerr << "Memory allocation failed" << std::endl;
			return NULL;
		}
		ipv4InstanceLocation->readMode = "mem";
		ipv4InstanceLocation->version = 4;
		ipv4InstanceLocation->fileName = fileName;
		if (ip_loadDat_v4(ipv4InstanceLocation, fileName) >= 0)
		{
			return ipv4InstanceLocation;
		}
		free(ipv4InstanceLocation);
		ipv4InstanceLocation = NULL;
		return NULL;
	}
	return NULL;
}

geo_ip *ipv4InstanceScenesFile(const char *fileName)
{
	geo_ip *ipv4InstanceScenes = (geo_ip *)malloc(sizeof(geo_ip));
	if (!ipv4InstanceScenes)
	{
		std::cerr << "Memory allocation failed" << std::endl;
		return NULL;
	}
	ipv4InstanceScenes->readMode = "file";
	ipv4InstanceScenes->version = 4;
	ipv4InstanceScenes->fileName = fileName;
	if (ip_loadDat_v4(ipv4InstanceScenes, fileName) >= 0)
	{
		return ipv4InstanceScenes;
	}
	free(ipv4InstanceScenes);
	ipv4InstanceScenes = NULL;
	return NULL;
}

geo_ip *ipv4InstanceScenesMemory(const char *fileName)
{
	static geo_ip *ipv4InstanceScenes = NULL;
	if (ipv4InstanceScenes == NULL)
	{
		ipv4InstanceScenes = (geo_ip *)malloc(sizeof(geo_ip));
		if (!ipv4InstanceScenes)
		{
			std::cerr << "Memory allocation failed" << std::endl;
			return NULL;
		}
		ipv4InstanceScenes->readMode = "mem";
		ipv4InstanceScenes->version = 4;
		ipv4InstanceScenes->fileName = fileName;
		if (ip_loadDat_v4(ipv4InstanceScenes, fileName) >= 0)
		{
			return ipv4InstanceScenes;
		}
		free(ipv4InstanceScenes);
		ipv4InstanceScenes = NULL;
		return NULL;
	}
	return NULL;
}

geo_ip *ipv4InstanceProxyFile(const char *fileName)
{
	geo_ip *ipv4InstanceScenes = (geo_ip *)malloc(sizeof(geo_ip));
	if (!ipv4InstanceScenes)
	{
		std::cerr << "Memory allocation failed" << std::endl;
		return NULL;
	}
	ipv4InstanceScenes->readMode = "file";
	ipv4InstanceScenes->version = 4;
	ipv4InstanceScenes->fileName = fileName;
	if (ip_loadDat_v4(ipv4InstanceScenes, fileName) >= 0)
	{
		return ipv4InstanceScenes;
	}
	free(ipv4InstanceScenes);
	ipv4InstanceScenes = NULL;
	return NULL;
}

geo_ip *ipv4InstanceProxyMemory(const char *fileName)
{
	static geo_ip *ipv4InstanceScenes = NULL;
	if (ipv4InstanceScenes == NULL)
	{
		ipv4InstanceScenes = (geo_ip *)malloc(sizeof(geo_ip));
		if (!ipv4InstanceScenes)
		{
			std::cerr << "Memory allocation failed" << std::endl;
			return NULL;
		}
		ipv4InstanceScenes->readMode = "mem";
		ipv4InstanceScenes->version = 4;
		ipv4InstanceScenes->fileName = fileName;
		if (ip_loadDat_v4(ipv4InstanceScenes, fileName) >= 0)
		{
			return ipv4InstanceScenes;
		}
		free(ipv4InstanceScenes);
		ipv4InstanceScenes = NULL;
		return NULL;
	}
	return NULL;
}

geo_ip *ipv6InstanceFile(const char *fileName)
{
	geo_ip *ipv6Instance = (geo_ip *)malloc(sizeof(geo_ip));
	if (!ipv6Instance)
	{
		std::cerr << "Memory allocation failed" << std::endl;
		return NULL;
	}
	ipv6Instance->readMode = "file";
	ipv6Instance->version = 6;
	ipv6Instance->fileName = fileName;
	if (ip_loadDat_v6(ipv6Instance, fileName) >= 0)
	{
		return ipv6Instance;
	}
	free(ipv6Instance);
	ipv6Instance = NULL;
	return NULL;
}

geo_ip *ipv6InstanceMemory(const char *fileName)
{
	static geo_ip *ipv6Instance = NULL;
	if (ipv6Instance == NULL)
	{
		ipv6Instance = (geo_ip *)malloc(sizeof(geo_ip));
		if (!ipv6Instance)
		{
			std::cerr << "Memory allocation failed" << std::endl;
			return NULL;
		}
		ipv6Instance->readMode = "mem";
		ipv6Instance->version = 6;
		ipv6Instance->fileName = fileName;
		if (ip_loadDat_v6(ipv6Instance, fileName) >= 0)
		{
			return ipv6Instance;
		}
		free(ipv6Instance);
		ipv6Instance = NULL;
		return NULL;
	}
	return NULL;
}

geo_ip *ip_instance(const char *fileName, const char *readMode, int version)
{
	static geo_ip *instance = NULL;
	if (instance == NULL)
	{
		geo_ip *ret = (geo_ip *)malloc(sizeof(geo_ip));
		if (!ret)
		{
			std::cerr << "Memory allocation failed" << std::endl;
			return NULL;
		}
		ret->readMode = readMode;
		ret->version = version;
		ret->fileName = fileName;
		if (version == 4)
		{
			if (ip_loadDat_v4(ret, fileName) >= 0)
			{
				return ret;
			}
			free(ret);
			ret = NULL;
			return NULL;
		}
		else if (version == 6)
		{

			if (ip_loadDat_v6(ret, fileName) >= 0)
			{
				return ret;
			}
			free(ret);
			ret = NULL;
			return NULL;
		}
		else
		{
			free(ret);
			ret = NULL;
			std::cout << "unsupported reading mode" << std::endl;
			return NULL;
		}
	}
	return instance;
}

int32_t ip_loadDat_v4(geo_ip *p, const char *fileName)
{

	long len = 0;
	int k, i, j;
	p->file = fopen(fileName, "rb");

	if (p->file == NULL)
	{
		printf("%s", "file open failed");
	}

	if (strcmp(p->readMode, "mem") == 0)
	{
		fseek(p->file, 0, SEEK_END);
		len = ftell(p->file);
		fseek(p->file, 0, SEEK_SET);
		p->buffer = (uint8_t *)malloc(len * sizeof(uint8_t));
		if (p->buffer == nullptr)
		{
			std::cerr << "create buffer failed" << std::endl;
			fclose(p->file);
			return -1;
		}
		if (fread(p->buffer, 1, len, p->file) != len)
		{
			std::cerr << "read file failed" << std::endl;
			fclose(p->file);
			return -1;
		}
		fclose(p->file);
	}
	for (k = 0; k < 256; k++)
	{
		i = k * 8 + 4;
		p->prefStart[k] = unpack_int_4byte(p, i);
		p->prefEnd[k] = unpack_int_4byte(p, i + 4);
	}

	return 0;
}

int32_t ip_loadDat_v6(geo_ip *p, const char *fileName)
{
	uint32_t len = 0;
	uint32_t k;
	int i, j;
	uint32_t RecordSize, offset, length;
	p->file = fopen(fileName, "rb");
	if (p->file == NULL)
	{
		printf("%s", "open file ");
		return -2;
	}

	if (strcmp(p->readMode, "mem") == 0)
	{
		fseek(p->file, 0, SEEK_END);
		len = ftell(p->file);
		fseek(p->file, 0, SEEK_SET);
		p->buffer = (uint8_t *)malloc(len * sizeof(uint8_t));
		if (fread(p->buffer, 1, len, p->file) != len)
		{
			std::cerr << "Error reading file" << std::endl;
			fclose(p->file);
			return -1;
		}
		fclose(p->file);
	}

	p->numbers = unpack_int_4byte(p, 4);
	memset(p->prefStart6, 0, sizeof(p->prefStart6));
	memset(p->prefEnd6, 0, sizeof(p->prefEnd6));
	for (k = 0; k < p->numbers; k++)
	{
		i = k * 12 + 4 + 4;
		p->prefStart6[unpack_int_4byte(p, i + 8)] = unpack_int_4byte(p, i);
		p->prefEnd6[unpack_int_4byte(p, i + 8)] = unpack_int_4byte(p, i + 4);
	}

	return 0;
}

std::vector<std::string> split(const std::string &s, char delimiter)
{
	std::vector<std::string> tokens;
	std::string token;
	std::istringstream tokenStream(s);
	while (std::getline(tokenStream, token, delimiter))
	{
		tokens.push_back(token);
	}
	return tokens;
}
long getPref(const char *ip)
{
	std::string str_ip = ip;
	int index = str_ip.find(":");
	str_ip = str_ip.substr(0, index);
	long pref = strtoll(str_ip.c_str(), NULL, 16);
	return pref;
}

std::string expandIPv6(const std::string &ipv6)
{
	std::vector<std::string> parts = split(ipv6, ':');
	std::string expandedIPv6;

	size_t emptyBlockIndex = ipv6.find("::");
	if (emptyBlockIndex != std::string::npos)
	{
		std::vector<std::string> firstHalf;
		std::vector<std::string> secondHalf;

		if (emptyBlockIndex > 0)
		{
			std::string firstPart = ipv6.substr(0, emptyBlockIndex);
			firstHalf = split(firstPart, ':');
		}
		if (emptyBlockIndex + 2 <= parts.size())
		{
			std::string secondPart = ipv6.substr(emptyBlockIndex + 2);
			secondHalf = split(secondPart, ':');
		}

		for (const auto &part : firstHalf)
		{
			if (!part.empty())
			{
				std::stringstream ss(part);
				uint16_t value;
				ss >> std::hex >> value;
				std::ostringstream formattedPart;
				formattedPart << std::setw(4) << std::setfill('0') << std::hex << value;
				expandedIPv6 += formattedPart.str() + ":";
			}
		}

		size_t zeroBlocks = 8 - firstHalf.size() - secondHalf.size();
		for (size_t i = 0; i < zeroBlocks; ++i)
		{
			expandedIPv6 += "0000:";
		}

		for (const auto &part : secondHalf)
		{
			if (!part.empty())
			{
				std::stringstream ss(part);
				uint16_t value;
				ss >> std::hex >> value;
				std::ostringstream formattedPart;
				formattedPart << std::setw(4) << std::setfill('0') << std::hex << value;
				expandedIPv6 += formattedPart.str() + ":";
			}
		}
	}
	else
	{
		expandedIPv6 = ipv6;
	}

	if (!expandedIPv6.empty() && expandedIPv6.back() == ':')
	{
		expandedIPv6.pop_back();
	}

	return expandedIPv6;
}

uint128_t ipv6ToInt(const std::string &ipv6)
{
	std::vector<uint16_t> parts;
	std::stringstream ss(ipv6);
	std::string part;
	while (std::getline(ss, part, ':'))
	{
		if (!part.empty())
		{
			uint16_t value;
			std::stringstream converter(part);
			converter >> std::hex >> value;
			parts.push_back(value);
		}
	}

	uint128_t result = 0;
	for (const auto &part : parts)
	{
		result = (result << 16) | part;
	}

	return result;
}
int deleteChar(char *s, char c)
{
	if (NULL == s)
	{
		return -1;
	}
	else
	{
		char *f = s;
		int i = 0, j = 0;
		while (*s)
		{
			i++;
			if (*s != c)
			{
				j++;
				*f = *s;
				f++;
			}
			s++;
		}
		*f = '\0';
		if (i == j)
			return 0;
		return i - j;
	}
}

char *ip_query(geo_ip *p, const char *ip)
{
	uint32_t pref, cur, low, high;
	if (p == NULL)
	{
		return NULL;
	}

	if (p->version == 4)
	{
		uint32_t intIP = ip_ip2long(p, ip, &pref);
		;
		low = p->prefStart[pref];
		high = p->prefEnd[pref];

		cur = (low == high) ? low : ip_binary_search(p, low, high, intIP);
		if (cur == 100000000)
		{
			const char *str = "|||||||||||||||||";
			char *nil = (char *)malloc(strlen(str) + 1);
			if (nil == NULL)
			{
				std::cerr << "Memory allocation failed" << std::endl;
				free(nil);
				return NULL;
			}
			strcpy(nil, str);
			return nil;
		}
		return get_addr(p, cur);
	}
	else if (p->version == 6)
	{
		pref = getPref(ip);
		low = p->prefStart6[pref];
		high = p->prefEnd6[pref];
		const uint128_t intIP = ipv6ToInt(expandIPv6(ip));
		cur = (low == high) ? low : ipv6_binary_search(p, low, high, intIP);
		return get_addr6(p, cur);
	}
	else
	{
		std::cerr << "unsupported reading mode" << std::endl;
		return NULL;
	}
}

long ip_binary_search(geo_ip *p, long low, long high, long k)
{
	long M = 0;
	while (low <= high)
	{
		long mid = (low + high) / 2;
		int j = 2052 + (mid * 9);
		uint32_t endipNum = unpack_int_4byte(p, j);

		if (endipNum >= k)
		{
			M = mid;
			if (mid == 0)
			{
				break;
			}
			high = mid - 1;
		}
		else
			low = mid + 1;
	}
	return M;
}

char *get_addr(geo_ip *p, uint32_t j)
{
	long index = 2052 + j * 9;

	uint32_t offset = unpack_int_4byte(p, index + 4);
	uint32_t length = unpack_int_1byte(p, index + 8);

	char *result = (char *)malloc((length + 1) * sizeof(char));
	if (result == NULL)
	{
		std::cerr << "Memory allocation failed" << std::endl;
		free(result);
		return NULL;
	}

	if (strcmp(p->readMode, "mem") == 0)
	{
		memcpy(result, p->buffer + offset, length);
		result[length] = '\0';
	}
	else if (strcmp(p->readMode, "file") == 0)
	{
		if (fseek(p->file, offset, SEEK_SET) != 0 || fread(result, 1, length, p->file) != length)
		{
			std::cerr << "Error reading file" << std::endl;
			free(result);
			return nullptr;
		}
		result[length] = '\0';
	}

	return result;
}
char *get_addr6(geo_ip *p, uint32_t j)
{
	long index = p->numbers * 12 + 4 + 4 + (j * 55);

	uint32_t length = unpack_int_1byte(p, index + 50 + 4);
	uint32_t offset = unpack_int_4byte(p, index + 50);
	char *result = (char *)malloc((length + 1) * sizeof(char));
	if (result == NULL)
	{
		std::cerr << "Memory allocation failed" << std::endl;
		return NULL;
	}

	if (strcmp(p->readMode, "mem") == 0)
	{
		memcpy(result, p->buffer + offset, length);
		result[length] = '\0';
	}
	else if (strcmp(p->readMode, "file") == 0)
	{
		if (fseek(p->file, offset, SEEK_SET) != 0 || fread(result, 1, length, p->file) != length)
		{
			std::cerr << "Error reading file" << std::endl;
			free(result);
			return NULL;
		}
		result[length] = '\0';
	}

	return result;
}

uint32_t ip_ip2long(geo_ip *p, const char *addr, uint32_t *prefix)
{
	uint32_t c, octet, t;
	uint32_t ipnum;
	int i = 3;

	octet = ipnum = 0;
	while ((c = *addr++))
	{
		if (c == '.')
		{
			ipnum <<= 8;
			ipnum += octet;
			i--;
			octet = 0;
		}
		else
		{
			t = octet;
			octet <<= 3;
			octet += t;
			octet += t;
			c -= '0';

			octet += c;
			if (i == 3)
			{
				*prefix = octet;
			}
		}
	}

	ipnum <<= 8;

	return ipnum + octet;
}

uint32_t unpack_int_4byte(geo_ip *p, long offset)
{
	if (strcmp(p->readMode, "mem") == 0)
	{
		return ((p->buffer[offset] & 0xff) | ((p->buffer[offset + 1] << 8) & 0xff00) | ((p->buffer[offset + 2] << 16) & 0xff0000) | ((p->buffer[offset + 3] << 24) & 0xff000000));
	}
	if (strcmp(p->readMode, "file") == 0)
	{
		char result[4];
		if (fseek(p->file, offset, SEEK_SET) != 0 || fread(result, 1, 4, p->file) != 4)
		{
			std::cerr << "Error reading file" << std::endl;
			fclose(p->file);
			return -1;
		}
		return (result[0] & 0xff) | ((result[1] << 8) & 0xff00) | ((result[2] << 16) & 0xff0000) | ((result[3] << 24) & 0xff000000);
	}
	std::cerr << "Error read mode" << std::endl;
	return -1;
}

uint32_t unpack_int_1byte(geo_ip *p, long offset)
{
	if (strcmp(p->readMode, "mem") == 0)
	{
		return (uint32_t)(p->buffer[offset] & 0xff);
	}
	if (strcmp(p->readMode, "file") == 0)
	{
		char result[1];
		if (fseek(p->file, offset, SEEK_SET) != 0 || fread(result, 1, 1, p->file) != 1)
		{
			std::cerr << "Error reading file" << std::endl;
			fclose(p->file);
			return -1;
		}
		return (uint32_t)(result[0] & 0xff);
	}
	std::cerr << "Error read mode" << std::endl;
	return -1;
}

long ipv6_binary_search(geo_ip *p, long low, long high, uint128_t k)
{
	long M = 0;
	char *resultEnd = (char *)malloc(51 * sizeof(char));
	if (resultEnd == NULL)
	{
		std::cerr << "Memory allocation failed" << std::endl;
	}
	while (low <= high)
	{
		long mid = (low + high) / 2;
		int j = p->numbers * 12 + 4 + 4 + (mid * 55);
		uint128_t endipNum;
		if (strcmp(p->readMode, "mem") == 0)
		{
			memcpy(resultEnd, p->buffer + j, 50);
			resultEnd[50] = '\0';
			deleteChar(resultEnd, '*');
			endipNum = uint128_t(resultEnd, 10);
		}
		else if (strcmp(p->readMode, "file") == 0)
		{
			if (fseek(p->file, j, SEEK_SET) != 0 || fread(resultEnd, 1, 50, p->file) != 50)
			{
				std::cerr << "Error reading file" << std::endl;
				fclose(p->file);
				free(resultEnd);
				resultEnd = NULL;
				return -1;
			}
			resultEnd[49] = '\0';
			deleteChar(resultEnd, '*');
			endipNum = uint128_t(resultEnd, 10);
		}
		else
		{
			std::cerr << "Error read mode" << std::endl;
			free(resultEnd);
			resultEnd = NULL;
			return -1;
		}

		if (endipNum >= k)
		{
			M = mid;
			if (mid == 0)
			{
				break;
			}
			high = mid - 1;
		}
		else
			low = mid + 1;
	}
	free(resultEnd);
	resultEnd = NULL;
	return M;
}
